import pandas as pd
from sklearn.model_selection import train_test_split

PREDICT_YEAR = '2023'
SKIP_PERIOD = 20 # same with boll period
DATA_TRAIN_COLUMN = ['open', 'high', 'low', 'close', 'vol', 'MA_5', 'MA_10', 'RSI', 'MEDIUM', 'STD', 'UPPER', 'LOWER']
DATA_TEST_COLUMN = ['DIRECTION']
PERFORMANCE_TRAIN_COLUMN = ['trade_date', 'open', 'high', 'low', 'close', 'vol', 'MA_5', 'MA_10', 'RSI', 'MEDIUM', 'STD', 'UPPER', 'LOWER']

def read_train_n_test_data(dir: str, test_size = 0.3) -> list[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    df_ori: pd.DataFrame = pd.read_csv(dir)
    assert(isinstance(df_ori, pd.DataFrame))
    df_ori = df_ori[~ df_ori['trade_date'].copy().astype(str).str.contains(PREDICT_YEAR)][SKIP_PERIOD:]
    factor = pd.DataFrame(df_ori[DATA_TRAIN_COLUMN]).reset_index()
    res = pd.DataFrame(df_ori[DATA_TEST_COLUMN])

    if test_size == 0:
        return factor, res['DIRECTION']
    return train_test_split(factor , res['DIRECTION'] , test_size=test_size)

def read_to_predict_data(dir: str) -> list[pd.DataFrame, pd.DataFrame]:
    df_ori: pd.DataFrame = pd.read_csv(dir)
    assert(isinstance(df_ori, pd.DataFrame))
    df_ori = df_ori[df_ori['trade_date'].copy().astype(str).str.contains(PREDICT_YEAR)]
    factor = pd.DataFrame(df_ori[PERFORMANCE_TRAIN_COLUMN]).reset_index()
    res = pd.DataFrame(df_ori[DATA_TEST_COLUMN])

    return factor, res['DIRECTION']


